from torch import nn
import torch.nn.functional as F


class ImageModel900(nn.Module):
    """
    处理900RGB图像
    """

    def __init__(self):
        super().__init__()
        # 输入n x (3x900x900)
        self.conv1 = nn.Conv2d(3, 8, 3)
        # 输出n x (8x898x898)
        self.mp1 = nn.MaxPool2d(4, stride=2)
        # 输出n x (8x448x448)
        self.conv2 = nn.Conv2d(8, 16, 3)
        # 输出n x (16x446x446)
        self.mp2 = nn.MaxPool2d(4, stride=2)
        # 输出n x (16x222x222)
        self.conv3 = nn.Conv2d(16, 32, 3, 3)
        # 输出n x (32x74x74)
        self.mp3 = nn.MaxPool2d(4, stride=2)
        # 输出n x (32x36x36)
        self.conv4 = nn.Conv2d(32, 32, 3, 3)
        # 输出n x (32x12x12)
        self.mp4 = nn.MaxPool2d(4, stride=4)
        # 输出n x (32x3x3)
        # BatchNorm1d对于batch大小敏感，所以这里就不用了
        # self.norm1 = nn.BatchNorm1d(288)
        self.fc1 = nn.Linear(288, 72)
        self.dp1 = nn.Dropout(p=0.4)
        # self.norm2 = nn.BatchNorm1d(72)
        self.fc2 = nn.Linear(72, 18)
        self.dp2 = nn.Dropout(p=0.4)
        # self.norm3 = nn.BatchNorm1d(18)
        self.fc3 = nn.Linear(18, 3)

    def forward(self, x):
        out = F.rrelu(self.conv1(x))
        out = self.mp1(out)
        out = F.rrelu(self.conv2(out))
        out = self.mp2(out)
        out = F.rrelu(self.conv3(out))
        out = self.mp3(out)
        out = F.rrelu(self.conv4(out))
        out = self.mp4(out)
        out = out.view((-1, 288))
        # out = self.norm1(out)
        out = self.fc1(out)
        out = self.dp1(out)
        out = F.rrelu(out)
        # out = self.norm2(out)
        out = self.fc2(out)
        out = self.dp2(out)
        out = F.rrelu(out)
        # out = self.norm3(out)
        out = self.fc3(out)
        # return F.softmax(out, dim=1)
        return out


if __name__ == "__main__":
    pass
